import os
import sys
from sklearn import feature_extraction
from sklearn.feature_extraction.text import TfidfTransformer
from sklearn.feature_extraction.text import CountVectorizer
import numpy as np


def compute_tf_idf(file):
    # 读取corpus
    with open(file, 'r', encoding='utf-8') as f:
        l = f.readlines()
        # 处理每条语句最后的换行符
        corpus = []
        for i in range(len(l)):
            corpus.append(l[i].strip())
        # 该类会将文本中的词语转换为词频矩阵，矩阵元素a[i][j] 表示j词在i类文本下的词频
        vectorizer = CountVectorizer()
        # 该类会统计每个词语的tf-idf权值
        transformer = TfidfTransformer()
        # 第一个fit_transform是计算tf-idf，第二个fit_transform是将文本转为词频矩阵
        tfidf = transformer.fit_transform(vectorizer.fit_transform(corpus))
        # 获取词袋模型中的所有词语 ,list
        word = vectorizer.get_feature_names()
        # 将tf-idf矩阵抽取出来，元素a[i][j]表示j词在i类文本中的tf-idf权重,  ndarray
        weight = tfidf.toarray()

        np.savez(r'F:\mypython\final_subject\bbdw\tf-idf-weight', weight)
        np_word = np.array(word)
        np.savez(r'F:\mypython\final_subject\bbdw\word', np_word)
        return weight


if __name__ == '__main__':
    w = compute_tf_idf('new_file.txt')
    print(w)